/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: dispatcher common api
 * Author: yanyefeng
 * Create: 2023-11-14
 */

#include "dispatcher_pub.h"
#include "dispatcher_ffts_pub.h"
#include "dispatcher_virtural_pub.h"
#include "dispatcher_aicpu_pub.h"
#include "dispatcher.h"
#include "externalinput_pub.h"

using namespace hccl;

typedef HcclResult (*FftsCounterCallBack)(const HcclDispatcher&, Stream &);
FftsCounterCallBack g_InitTaskCallback = nullptr;
FftsCounterCallBack g_LaunchTaskCallback = nullptr;
void RegisterInitTaskCallBack(HcclResult (*p1)(const HcclDispatcher &, hccl::Stream &))
{
    g_InitTaskCallback = p1;
}

void RegisterLaunchTaskCallBack(HcclResult (*p1)(const HcclDispatcher &, hccl::Stream &))
{
    g_LaunchTaskCallback = p1;
}

HcclResult RegisterLoadTaskCallBack(HcclDispatcher dispatcherPtr, void *userPtr,
    void (*p1)(void *userPtr, void *param, u32 length))
 
{
    CHK_PTR_NULL(dispatcherPtr);
    reinterpret_cast<DispatcherPub*>(dispatcherPtr)->RegLoadTaskCallBack(userPtr, p1);
    return HCCL_SUCCESS;
}

HcclResult HcclDispatcherInit(DispatcherType type, const s32 devicePhyId, HcclDispatcher *dispatcher)
{
    CHK_RET(DlProfFunc::GetInstance().DlProfFunctionInit());
    CHK_PTR_NULL(dispatcher);
    u32 deviceLogicId = INVALID_UINT;
    if (static_cast<s32>(devicePhyId) != HOST_DEVICE_ID) {
        CHK_RET(hrtGetDeviceIndexByPhyId(devicePhyId, deviceLogicId));
    } else {
        deviceLogicId = devicePhyId;
    }

    DispatcherPub *pDispatcher = nullptr;
    if (type == DispatcherType::DISPATCHER_NORMAL) {
        if (GetExternalInputHcclEnableFfts()) {
            pDispatcher = new (std::nothrow) DispatcherFFTS(deviceLogicId);
        } else {
            pDispatcher = new (std::nothrow) DispatcherPub(deviceLogicId);
        }
    } else if (type == DispatcherType::DISPATCHER_VIRTURAL) {
        pDispatcher = new (std::nothrow) DispatcherVirtural(deviceLogicId);
    } else {
        HCCL_ERROR("Not support the dispatcher type[%d]", type);
        return HCCL_E_NOT_SUPPORT;
    }
    CHK_PTR_NULL(pDispatcher);
    HcclResult ret = pDispatcher->Init();
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("Dispatcher init failed, type[%d]", type);
        delete pDispatcher;
        pDispatcher = nullptr;
        return ret;
    }
    *dispatcher = pDispatcher;
    return HCCL_SUCCESS;
}

HcclResult HcclDispatcherDestroy(HcclDispatcher dispatcherPtr)
{
    if (dispatcherPtr != nullptr) {
        DispatcherPub* dispatcher = reinterpret_cast<DispatcherPub*>(dispatcherPtr);
        delete dispatcher;
        dispatcherPtr = nullptr;
    }
    return HCCL_SUCCESS;
}

    HcclResult HcclSetGlobalWorkSpace(HcclDispatcher dispatcherPtr, std::vector<void *> &globalWorkSpaceAddr)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SetGlobalWorkSpace(globalWorkSpaceAddr);
}

HcclResult HcclSetQosCfg(HcclDispatcher dispatcherPtr, const u32 qosCfg)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SetQosCfg(qosCfg);
}

HcclResult HcclResetQosCfg(HcclDispatcher dispatcherPtr)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->ResetQosCfg();
}

HcclResult HcclGetQosCfg(HcclDispatcher dispatcherPtr, u32 *qosCfg)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(qosCfg);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->GetQosCfg(*qosCfg);
}

HcclResult HcclSetNotifyWaitMode(HcclDispatcher dispatcherPtr, const SyncMode notifyWaitMode)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SetNotifyWaitMode(notifyWaitMode);
}

HcclResult HcclGetNotifyWaitMode(HcclDispatcher dispatcherPtr, SyncMode *notifyWaitMode)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(notifyWaitMode);

    *notifyWaitMode = reinterpret_cast<DispatcherPub*>(dispatcherPtr)->GetNotifyWaitMode();
    return HCCL_SUCCESS;
}

HcclResult HcclD2DMemcpyAsync(HcclDispatcher dispatcherPtr, DeviceMem &dst, const DeviceMem &src,
    Stream &stream, const u32 remoteUserRank, const LinkType linkType)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(dst.ptr());
    CHK_PTR_NULL(src.ptr());

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->MemcpyAsync(dst, src, stream, remoteUserRank, linkType);
}

HcclResult HcclMemcpyAsync(HcclDispatcher dispatcherPtr, void *dst, const uint64_t destMax, const void *src,
    uint64_t count, const HcclRtMemcpyKind kind, Stream &stream, const u32 remoteUserRank,
    const LinkType linkType)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(dst);
    CHK_PTR_NULL(src);

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->MemcpyAsync(dst, destMax, src, count, kind,
        stream, remoteUserRank, linkType);
}

HcclResult HcclReduceAsync(HcclDispatcher dispatcherPtr, void *src, uint64_t count, const HcclDataType datatype,
    const HcclReduceOp reduceOp, Stream &stream, void *dst, const u32 remoteUserRank,
    const LinkType linkType, const u64 reduceAttr)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(dst);
    CHK_PTR_NULL(src);

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->ReduceAsync(src, count, datatype, reduceOp,
        stream, dst, remoteUserRank, linkType, reduceAttr);
}

HcclResult HcclDispatcherWaitValue(HcclDispatcher dispatcherPtr, hccl::Stream &stream, u64 waitAddr, u64 valueAddr, bool reset)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->WaitValue(stream, waitAddr, valueAddr, reset);
}
HcclResult HcclDispatcherWriteValue(HcclDispatcher dispatcherPtr, hccl::Stream &stream, u64 writeAddr, u64 valueAddr)
{
    CHK_PTR_NULL(dispatcherPtr);
    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->WriteValue(stream, writeAddr, valueAddr);
}

HcclResult HcclSignalRecord(HcclDispatcher dispatcherPtr, HcclRtNotify signal, Stream &stream, u32 userRank,
    u64 offset, s32 stage, bool inchip, u64 signalAddr)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(signal);

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SignalRecord(signal, stream, userRank,
        offset, stage, inchip, signalAddr);
}
HcclResult HcclSignalWait(HcclDispatcher dispatcherPtr, HcclRtNotify signal, Stream &stream, u32 userRank,
    u32 remoteUserRank, s32 stage, bool inchip)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(signal);

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SignalWait(signal, stream,
        userRank, remoteUserRank, stage, inchip);
}

HcclResult LaunchTask(HcclDispatcher dispatcherPtr, Stream &stream)
{
    CHK_PTR_NULL(dispatcherPtr);
    if (g_LaunchTaskCallback != nullptr) {
        CHK_RET(g_LaunchTaskCallback(dispatcherPtr, stream));
    }
    std::vector<Stream> subStreams;
    HcclResult ret = reinterpret_cast<DispatcherPub*>(dispatcherPtr)->LaunchTasksEx(stream, subStreams);
    return ret;
}

HcclResult LaunchTaskExtend(HcclDispatcher dispatcherPtr, Stream &stream, std::vector<Stream> &subStreams)
{
    CHK_PTR_NULL(dispatcherPtr);
    if (g_LaunchTaskCallback != nullptr) {
        CHK_RET(g_LaunchTaskCallback(dispatcherPtr, stream));
    }
    return reinterpret_cast<DispatcherPub *>(dispatcherPtr)->LaunchTasksEx(stream, subStreams);
}

HcclResult InitTask(HcclDispatcher dispatcherPtr, hccl::Stream &stream, const bool enableCache, const std::string &key)
{
    CHK_PTR_NULL(dispatcherPtr);

    CHK_RET(reinterpret_cast<DispatcherPub*>(dispatcherPtr)->ResetFftsCtx(enableCache, key));
    if (g_InitTaskCallback != nullptr) {
        CHK_RET(g_InitTaskCallback(dispatcherPtr, stream));
    }
    return HCCL_SUCCESS;
}

HcclResult AddRetryPreamble(HcclDispatcher dispatcherPtr, hccl::Stream &stream)
{
    CHK_PTR_NULL(dispatcherPtr);

    CHK_RET(reinterpret_cast<DispatcherPub*>(dispatcherPtr)->AddRetryPreamble(stream));
    return HCCL_SUCCESS;
}

HcclResult SetNormalMode(HcclDispatcher dispatcherPtr)
{
    CHK_PTR_NULL(dispatcherPtr);

    reinterpret_cast<DispatcherPub*>(dispatcherPtr)->SetNormalMode();

    return HCCL_SUCCESS;
}

HcclResult IsCtxInitialized(HcclDispatcher dispatcherPtr, bool *fftsCtxInitFlag)
{
    CHK_PTR_NULL(dispatcherPtr);
    CHK_PTR_NULL(fftsCtxInitFlag);

    reinterpret_cast<DispatcherPub*>(dispatcherPtr)->JudgeFftsCtxInitialized(*fftsCtxInitFlag);
    return HCCL_SUCCESS;
}

HcclResult HcclGetCallbackResult(HcclDispatcher dispatcherPtr)
{
    CHK_PTR_NULL(dispatcherPtr);

    return reinterpret_cast<DispatcherPub*>(dispatcherPtr)->GetCallbackResult();
}